from operator import itemgetter

import dotenv
from langchain_community.chat_models import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel

# 通过runnable parallel 实现一个检索过程，将检索的内容，传递给chain 就是整个聊天过程
# chain = "12345"|prompt|llm|parser 前面的“123456”并不能并入到主流程

dotenv.load_dotenv()

def retriver_from_anywhere(query:str):
    print(f"开始检索问题{query}的相关文档")
    return "我叫博小瑞，今年11岁，存款110块"

# 构建一个prompt
prompt = ChatPromptTemplate.from_template("""
请根据用户的问题来回答，可以参考上下文进行问题的回答，
<context>
{context}
</context>
用户的问题是：
{query}
""")
# llm的构建
llm = ChatOpenAI(model="gpt-3.5-turbo-16k")
# 输出解析器
paraser = StrOutputParser()

chain = RunnableParallel({
    "context": lambda x:retriver_from_anywhere(x["query"]),
    "query":itemgetter("query")
})|prompt|llm|paraser

print(chain.invoke({"query":"你好，我叫什么名字？"}))
print(chain.invoke({"query":"你好，我有多少钱？"}))
print(chain.invoke({"query":"你好，我多少岁？"}))
print(chain.invoke({"query":"你好，我是谁？"}))


# print(chain.invoke({"query":"你好，我叫什么名？"}))
# print